import logging
import scipy.misc
from scipy.ndimage.filters import median_filter
# import matplotlib.pyplot as plt
import numpy as np
from scipy.ndimage.morphology import binary_opening, binary_closing
import pbcvt
import xml.etree.ElementTree as ET
import os
import scipy.ndimage
import cv2
from  os import mkdir
from  os.path import exists
import csv
import math
from logging import warning, info, debug
import random

# import caffe

HOME_DIR = os.environ['HOME']
DB_ROOT = HOME_DIR + '/data/ddb1_v02_01/'
TRAIN_SET_FILE = HOME_DIR + '/data/ddb1_v02_01/ddb1_v02_01_train_plain.txt'
TEST_SET_FILE = HOME_DIR + '/data/ddb1_v02_01/ddb1_v02_01_test_plain.txt'
TMP_OUT = 'out_dir/'


def imread(imgname, root=DB_ROOT):
    """Read image file form predifined file, To load different
    database, just need to change this function

    Args:
        imagename: name of the file, not the whole path
    Returns:
        Image read from the given name
    """
    img = scipy.misc.imread(root + imgname)
    # size = np.floor(np.array(img.shape) * 0.7).astype(np.int32)
    # img = scipy.misc.imresize(img, size)
    img = img / 255
    return img


def imsave(imagename, img, scale=False, prefix=TMP_OUT):
    """Save image file to folder defined in TMP_OUT
    Args:
        imagename: the output file name
        img: byte array of image
    Return:
        None
    """
    assert type(imagename) == str
    logging.debug('save to ' + imagename)
    if not scale:
        img[img < 0] = 0
        img[img > 1] = 1
    fname = prefix + imagename + '.png'
    scipy.misc.imsave(fname, img)
    return fname


def vessel_seg(image):
    return pbcvt.vesselseg(image) > 0
    # return image


def imsaveWithIndex(imagename, img, index=-1, scale=True, prefix=TMP_OUT):
    if index > -1:
        return imsave('%05d_%s' % (index, imagename), img, scale, prefix)
    else:
        raise 'err index ' + str(index)


        # def histgram(image):
        #     plt.hist(image)
        #     plt.title("Gaussian Histogram")
        #     plt.xlabel("Value")
        #     plt.ylabel("Frequency")
        #     fig = plt.gcf()
        #     plt.show()
        # plot_url = py.plot_mpl(fig, filename='mpl-basic-histogram')
        # print(plot_url)


def candateSelection(image, tmpFileIndex=-1):
    """Generate candates of MA
    Args:
        image: image of fundus image
        tmpFileIndex: if larger than -1, than write tmp output image to file
    Returns:
        A list of tuples containing (row,col,radius) of MA center
    """
    invGreen = 1 - image[:, :, 1]
    imsaveWithIndex('0invGreen', invGreen, tmpFileIndex)
    bg = median_filter(invGreen, (25, 25), mode='nearest')
    imsaveWithIndex('1bg', bg, tmpFileIndex)
    sc = invGreen - bg
    imsaveWithIndex('2sc', sc, tmpFileIndex)
    cand = (sc > 0.01) + 0
    imsaveWithIndex('3cand', cand, tmpFileIndex)
    scale = 700 / image.shape[0];
    vessel = 1 - vessel_seg(scipy.misc.imresize(image, scale))
    # logging.info(vessel.max())
    vessel = scipy.misc.imresize(vessel, image.shape)
    imsaveWithIndex('4vessel', vessel, tmpFileIndex)
    # remove vessel
    vesselRemovel = cand * vessel
    imsaveWithIndex('5cand_remove_vessel', vesselRemovel, tmpFileIndex)
    result = binary_closing(vesselRemovel, np.ones((3, 3)))
    result = binary_opening(result, np.ones((3, 3)))
    imsaveWithIndex('6cand_opening', result, tmpFileIndex)
    return result


def extractGroundTruth(xmlFileList):
    cord = []
    for file in xmlFileList:
        root = ET.parse(DB_ROOT + file)
        markinglist = root.find('markinglist')
        for marking in markinglist.findall('marking'):
            markingtype = marking.find('markingtype')
            circleregion = marking.find('circleregion')
            if markingtype.text == 'Red_small_dots':
                centroid = marking.find('.//centroid').find('coords2d')
                cordtext = centroid.text.split(',')
                cord.append((int(cordtext[1]), int(cordtext[0])))
    return cord


def drawPoint(img, cord, length=20, shape='cross'):
    for c in cord:
        crow = c[0]
        ccol = c[1]
        for row in range(max(0, crow - length), min(img.shape[0], crow + length)):
            for col in range(max(0, ccol - length), min(ccol + length, img.shape[1])):
                dist = abs(ccol - col) + abs(crow - row)
                if not (row == crow or col == ccol):
                    continue
                if dist < 10:
                    continue
                img[row][col] = 1
    return img


_DFS_PATH = [
    # 4-connectivity
    (0, 1), (1, 0), (0, -1), (-1, 0),
    # 8-connectivity
    (1, 1), (1, -1), (-1, 1), (-1, -1)
]


def minmax(arr, val):
    if arr[0] == -1:
        return val, val
    return min(val, arr[0]), max(val, arr[1])


def evaluate_result(segment, coord):
    TP = []
    threshold_len_sq = 10 ** 2
    for tr, tc in coord:
        selected = False
        min_sqr_err = 1000000
        minval = (-1, -1)
        for ar, ac in segment:
            sq_err = (ar - tr) ** 2 + (ac - tc) ** 2
            if sq_err < min_sqr_err:
                minval = (ar, ac)
                min_sqr_err = sq_err
        if min_sqr_err < threshold_len_sq:
            TP.append(minval)

    FP = segment[:-len(TP)]
    return TP, FP


def gen_mask(image):
    """
    :param image: the original fundus image with 3-channel
    :return: the binary mask of fundus image
    """
    v = np.zeros(image.shape[:-1])
    for i in range(3):
        v = v + image[:, :, i]
    v = v > 0.05
    return scipy.ndimage.morphology.binary_opening(v)


def get_local_max(image):
    """
    Get local max of the image
    :param image: grave level image
    :return: An array containing array of [row, col, label]
    """
    mask = gen_mask(image)
    # imsaveWithIndex('mask', mask)
    img = 1 - image[:, :, 1]
    img = scipy.ndimage.filters.gaussian_filter(img, 2)
    maxImg = scipy.ndimage.filters.maximum_filter(img, (5, 5))
    localmax = (maxImg == img) + 0
    localmax = localmax * mask
    retval, labels, stats, centroids = cv2.connectedComponentsWithStats(localmax.astype(np.uint8))
    for ind in range(len(centroids)):
        centroids[ind] = [centroids[ind, 1], centroids[ind, 0]]
    # logging.info(len(centroids))
    index = 0
    retval = []
    for x in centroids:
        retval.append([x[0], x[1], index, stats[index, cv2.CC_STAT_AREA]])
        index += 1
    peak = pbcvt.findPeak((img * 255).astype(np.uint8), np.array(retval).astype(np.int32))
    if peak is None:
        return []
    peak = peak[:, 0:2]
    debug(len(peak))
    debug('peak finished')
    return peak


class DDBDatabase:
    def __init__(self, file_list):
        self.image_file = []
        self.annotation = []
        with open(file_list, 'r') as csvfile:
            spreader = csv.reader(csvfile, delimiter=' ')
            for row in spreader:
                self.image_file.append(row[0])
                self.annotation.append(row[1:])
        self.current_index = 0

    def get_data(self, index=-1):
        if index == -1:
            index = self.current_index
            self.current_index += 1
        if index >= len(self.image_file):
            return None, None
        return scipy.misc.imread(DB_ROOT + self.image_file[index]), extractGroundTruth(self.annotation[index])


def main(imgfilename, annotations, index):
    index = int(index)
    img = imread(imgfilename)
    cord = extractGroundTruth(annotations)
    peak = get_local_max(img)
    TP, TF = evaluate_result(peak, cord)
    logging.info('%d\t%d\t%d' % (len(TP), len(TF), len(cord)))
    # --------------------------------------


def generate_dataset(train_dir='data.train/', test_dir='data.test/', size=85):
    if not exists(train_dir):
        mkdir(train_dir)
    if not exists(test_dir):
        mkdir(test_dir)
    index = 0
    half_width = math.floor(size / 2)
    for db_root in [TRAIN_SET_FILE, TEST_SET_FILE]:
        dd = DDBDatabase(db_root)
        if db_root.find('test') >= 0:
            odir = train_dir
        else:
            odir = test_dir
        listfile = open(odir + 'filelist.txt', 'w')
        print('Phase', db_root)
        while True:
            img, annotation = dd.get_data()
            counter = 0
            if img is None:
                break
            for row, col in annotation:
                if not ((0 <= row - half_width < row + half_width < img.shape[0]) and
                            (0 <= col - half_width < col + half_width < img.shape[0])):
                    # warning('index ' + str((row, col)) + ' out of range')
                    continue
                print('handlng %10d : %10d\r' % (dd.current_index, index), end='')
                fname = imsaveWithIndex('1img',
                                        img[row - half_width:row + half_width, col - half_width:col + half_width + 1,
                                        :], index, prefix=odir)
                index += 1
                counter += 1
                # double the true samples
                listfile.write('%s 1\n' % fname)
                listfile.write('%s 1\n' % fname)
                listfile.write('%s 1\n' % fname)
            ff = get_local_max(img / 255)
            for _ in range(9 * counter):
                while True:
                    xx = ff[random.randint(0,len(ff)-1)]
                    row = int(xx[0])
                    col = int(xx[1])
                    if not ((0 <= row - half_width < row + half_width < img.shape[0]) and
                                (0 <= col - half_width < col + half_width < img.shape[0])):
                        # warning('index ' + str((row, col)) + ' out of range')
                        continue
                    valid = True
                    for trow,tcol in annotation:
                        if (trow-row)**2+(tcol-col)**2 <= 50:
                            valid = False
                            break
                    if not valid:
                        continue
                    fname = imsaveWithIndex('0img',
                                            img[row - half_width:row + half_width, col - half_width:col + half_width + 1,
                                            :], index, prefix=odir)
                    index += 1
                    listfile.write('%s 0\n' % fname)
                    break
        listfile.close()


if __name__ == '__main__':
    import sys
    from logging import debug

    FORMAT = '[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
    DATEFORMAT = '%H:%M:%S'
    logging.basicConfig(level=logging.INFO, format=FORMAT, datefmt=DATEFORMAT)
    logging.debug('start')

    if (sys.argv[1] == 'gendata'):
        generate_dataset()
    else:
        x = [[1], 2, 3]
        for i in x:
            i = 0
        debug(x)
        debug('image name ' + str(sys.argv[1]))
        debug('annotations ' + str(sys.argv[2:-1]))
        debug('index ' + str(sys.argv[-1]))
        main(sys.argv[1], sys.argv[2:-1], sys.argv[-1])
